package com.ai.flight.ai.config;

import java.io.File;
import java.io.FileInputStream;
import java.io.InputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.List;

import com.ai.AiApplication;
import com.ai.flight.ai.service.LoggingAdvisor;
import com.ai.flight.ai.tools.FlightBookingAiTools;
import com.ai.flight.ai.tools.WeatherTools;
import com.alibaba.cloud.ai.advisor.DocumentRetrievalAdvisor;
import com.alibaba.cloud.ai.advisor.RetrievalRerankAdvisor;
import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel;
import com.alibaba.cloud.ai.dashscope.rerank.DashScopeRerankModel;
import com.alibaba.cloud.ai.document.TextDocumentParser;
import com.alibaba.cloud.ai.memory.redis.RedisChatMemoryRepository;
import org.elasticsearch.client.RestClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStoreOptions;
import org.springframework.ai.vectorstore.elasticsearch.SimilarityFunction;
import org.springframework.ai.vectorstore.elasticsearch.autoconfigure.ElasticsearchVectorStoreProperties;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.CommandLineRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;

/**
 * @Author: suren@517na.com
 * @CreateTime: 2025-06-02
 * @Description: 航班智能客户配置类
 */

@Configuration
public class AiConfig {

    private final Logger logger = LoggerFactory.getLogger(AiConfig.class);

    @Value("${ali.api-key}")
    private String apyKey;

    @Value("${spring.ai.memory.redis.port}")
    private Integer redisPort;

    @Value("${spring.ai.memory.redis.host}")
    private String redisHost;

    @Value("${spring.ai.memory.redis.password}")
    private String redisPassword;

    @Value("${spring.ai.memory.redis.timeout}")
    private int redisTimeout;

    @Autowired
    private FlightBookingAiTools flightBookingAiTools;
    @Autowired
    private WeatherTools weatherTools;

    @Bean
    public ChatClient qwenChatClient(DashScopeChatModel dashScopeChatModel,
                                     ChatMemory chatMemory,
                                     DocumentRetriever vectorStoreDocumentRetriever,
                                     ElasticsearchVectorStore elasticsearchVectorStore) {
        return ChatClient.builder(dashScopeChatModel)
                //预设角色
                .defaultSystem("""
                       您是“Funnair”航空公司的客户聊天支持代理。请以友好、乐于助人且愉快的方式来回复。
                       您正在通过在线聊天系统与客户互动。
                       在提供有关预订或取消预订的信息之前，您必须始终从用户处获取以下信息：预订号、客户姓名。
                       在询问用户之前，请检查消息历史记录以获取此信息。
                       在更改预订之前，您必须确保条款允许这样做。
                       如果更改需要收费，您必须在继续之前征得用户同意。
                       使用提供的功能获取预订详细信息、更改预订和取消预订。
                       如果需要，可以调用相应函数调用完成辅助动作。
                       请讲中文。
                       今天的日期是 {current_date}.
                        """)
                .defaultAdvisors(
                        MessageChatMemoryAdvisor.builder(chatMemory).build(),//对话记忆拦截器，用于存储历史对话,
                        new LoggingAdvisor(),
//                        RetrievalAugmentationAdvisor.builder()
//                                .documentRetriever(vectorStoreDocumentRetriever)
//                                .build()
                        //使用下面这种方式（RetrievalRerankAdvisor 和 DocumentRetrievalAdvisor配合使用）才能使Rag生效，上面注释掉的拦截器Rag不生效
                        new RetrievalRerankAdvisor(elasticsearchVectorStore,
                                new DashScopeRerankModel(
                                        DashScopeApi.builder().apiKey(apyKey).build()
                                )
                        ),
                        new DocumentRetrievalAdvisor(vectorStoreDocumentRetriever)
                )
                .defaultTools(flightBookingAiTools, weatherTools)
                .build();
    }

    @Bean
    public ChatMemory chatMemory(RedisChatMemoryRepository redisChatMemoryRepository) {
        return MessageWindowChatMemory.builder()
                .chatMemoryRepository(redisChatMemoryRepository)
                .maxMessages(20)
                .build();
    }


    /**
     * 内容检索器
     * @return bean
     */
    @Bean
    public DocumentRetriever vectorStoreDocumentRetriever(ElasticsearchVectorStore elasticsearchVectorStore) {
        return VectorStoreDocumentRetriever.builder()
                .vectorStore(elasticsearchVectorStore)
                .similarityThreshold(0.7) //相似度阈值
                .topK(2)       //返回top2的结果
                .build();
    }

    /**
     * Redis 持久化对话 Bean
     * @return bean
     */
    @Bean
    public RedisChatMemoryRepository redisChatMemoryRepository() {
        return RedisChatMemoryRepository
                .builder()
                .port(this.redisPort)
                .host(this.redisHost)
                .password(this.redisPassword)
                .timeout(this.redisTimeout)
                .build();
    }


    /**
     * 向量存储
     *
     * @param dashScopeEmbeddingModel 前文向量模型
     * @return 存储容器
     */
    @Bean
    public ElasticsearchVectorStore elasticsearchVectorStore(RestClient restClient,
                                                             DashScopeEmbeddingModel dashScopeEmbeddingModel,
                                                             ElasticsearchVectorStoreProperties esOptions) {
//        return SimpleVectorStore.builder(dashScopeEmbeddingModel).build(); //内存存储向量
        ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
        options.setDimensions(esOptions.getDimensions());
        options.setIndexName(esOptions.getIndexName());
        options.setSimilarity(esOptions.getSimilarity());
        return ElasticsearchVectorStore.builder(restClient, dashScopeEmbeddingModel)
                .options(options)
                .initializeSchema(esOptions.isInitializeSchema())
                .build(); //使用ES作为向量数据库
    }


    /**
     * 导入知识库数据到向量容器中
     * @param elasticsearchVectorStore
     * @return
     * @throws URISyntaxException
     */
    @Bean
    CommandLineRunner importDataToVectorStore(ElasticsearchVectorStore elasticsearchVectorStore) throws URISyntaxException {
//        URI uri = AiApplication.class.getClassLoader().getResource("rag/terms-of-service.txt").toURI();
//        File file = new File(uri);
//        return args -> {
//            TextDocumentParser textDocumentParser = new TextDocumentParser();
//            InputStream inputStream = new FileInputStream(file);
//            List<Document> documents = textDocumentParser.parse(inputStream);
//
//            List<Document> documentList = new TokenTextSplitter().apply(documents);
//            elasticsearchVectorStore.add(documentList);
//
//        };
        return args -> {

            // 转换成document
            PagePdfDocumentReader reader = new PagePdfDocumentReader(new ClassPathResource("rag/龙玉虎的简历.pdf"));
            List<Document> documents = reader.get();

            logger.info("{} document loaded", documents.size() );

            //切割成 trunks
            List<Document> splitDocuments = new TokenTextSplitter().apply(documents);
            logger.info("{} documents split", splitDocuments.size());

            //向量存储到es中
            logger.info("create embedding and save to ES");

            elasticsearchVectorStore.add(splitDocuments);
        };
    }
}
